Skip to content

Conversation

@glwagner
Copy link
Collaborator

@glwagner glwagner commented Nov 12, 2025

This PR starts adding some integration tests. I noticed that the first thing I tried failed, so we'll have to work our way up.

Working on this with @giordano and @jtbuch

@github-actions
Copy link
Contributor

github-actions bot commented Nov 12, 2025

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.
diff --git a/test/integration/Oceananigans/runtests.jl b/test/integration/Oceananigans/runtests.jl
index dab12a36..41a64e6e 100644
--- a/test/integration/Oceananigans/runtests.jl
+++ b/test/integration/Oceananigans/runtests.jl
@@ -3,14 +3,14 @@ using GPUArraysCore: @allowscalar
 using Oceananigans
 using Test
 
-function kinetic_energy(model, ν, K, dt=1)
+function kinetic_energy(model, ν, K, dt = 1)
     vitd = VerticallyImplicitTimeDiscretization()
     new_closure = ScalarDiffusivity(vitd; ν)
     tracer_names = keys(model.tracers)
     new_closure = Oceananigans.TurbulenceClosures.with_tracers(tracer_names, new_closure)
     model.closure = new_closure
 
-    for n = 1:10
+    for n in 1:10
         time_step!(model, dt)
     end
 
@@ -20,17 +20,19 @@ function kinetic_energy(model, ν, K, dt=1)
 end
 
 @testset "Column model with ScalarDiffusivity" begin
-    ν = 1e-3
-    grid = RectilinearGrid(size=128, z=(-64, 64), topology=(Flat, Flat, Bounded))
+    ν = 1.0e-3
+    grid = RectilinearGrid(size = 128, z = (-64, 64), topology = (Flat, Flat, Bounded))
     vitd = VerticallyImplicitTimeDiscretization()
     closure = ScalarDiffusivity(vitd; ν)
-    m = HydrostaticFreeSurfaceModel(; grid, closure, coriolis=FPlane(f=1e-4),
-        tracers=:b, buoyancy=BuoyancyTracer())
+    m = HydrostaticFreeSurfaceModel(;
+        grid, closure, coriolis = FPlane(f = 1.0e-4),
+        tracers = :b, buoyancy = BuoyancyTracer()
+    )
 
-    N² = 1e-6
+    N² = 1.0e-6
     bᵢ(z) = N² * z
     uᵢ(z) = exp(-z^2 / 10)
-    set!(m, b=bᵢ, u=uᵢ)
+    set!(m, b = bᵢ, u = uᵢ)
 
     u, v, w = m.velocities
     ke = (u^2 + v^2) / 2
@@ -39,11 +41,13 @@ end
     dm = Enzyme.make_zero(m)
     dK = Enzyme.make_zero(K)
 
-    dKdν = autodiff(set_strong_zero(Enzyme.ReverseWithPrimal),
-                    kinetic_energy, Active,
-                    Duplicated(m, dm),
-                    Active(ν),
-                    Duplicated(K, dK))
+    dKdν = autodiff(
+        set_strong_zero(Enzyme.ReverseWithPrimal),
+        kinetic_energy, Active,
+        Duplicated(m, dm),
+        Active(ν),
+        Duplicated(K, dK)
+    )
 
     @test dKdν[1][1] != 0
 end

@glwagner
Copy link
Collaborator Author

Looks like I need to be a little less ambitious probably


compute!(K)

return @allowscalar first(K)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@giordano @glwagner so this is failing atm because allowscalar isn't currently supported in enzyme.jl natively [we should add cc @vchuravy ]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants